package main

import (
	"archive/zip"
	"flag"
	"fmt"
	"github.com/sabhiram/go-gitignore"
	"github.com/schollz/progressbar/v3"
	yekazip "github.com/yeka/zip"
	"io"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/text/encoding/simplifiedchinese"
	"golang.org/x/text/transform"
)

var ignoreFlag string
var password string

func init() {
	flag.StringVar(&ignoreFlag, "ignore", "", "额外的忽略规则（逗号分隔），例如：\"*.log,tmp/,secret.txt\"")
	flag.StringVar(&password, "password", "", "压缩密码")
	flag.Parse()
}

func main() {
	args := flag.Args()
	if len(args) < 1 {
		printUsage()
		return
	}
	srcDir, outputZip := processArgs(args)

	ign := loadIgnore(srcDir, ignoreFlag)

	basePath := filepath.Clean(srcDir)
	filesToZip, err := collectFiles(basePath, ign)
	if err != nil {
		fmt.Println("扫描文件失败:", err)
		return
	}

	if err := createZipArchive(filesToZip, basePath, outputZip, password); err != nil {
		fmt.Println("\n压缩失败:", err)
		return
	}
	fmt.Println("\n✅ 压缩完成:", outputZip)
}

func createZipArchive(filesToZip []string, basePath string, outputZip string, password string) error {
	bar := progressbar.Default(int64(len(filesToZip)))

	zipFile, err := os.Create(outputZip)
	if err != nil {
		return fmt.Errorf("无法创建zip文件: %w", err)
	}
	defer zipFile.Close()

	if password == "" {
		zipWriter := zip.NewWriter(zipFile)
		defer zipWriter.Close()

		for _, relPath := range filesToZip {
			fullPath := filepath.Join(basePath, relPath)
			file, err := os.Open(fullPath)
			if err != nil {
				fmt.Printf("\n无法打开: %s\n%s\n", relPath, err)
				continue
			}

			info, _ := file.Stat()
			header, err := zip.FileInfoHeader(info)
			if err != nil {
				file.Close()
				fmt.Printf("\n获取文件头失败: %s\n%s\n", relPath, err)
				continue
			}
			header.Name = relPath
			header.Method = zip.Deflate

			writer, err := zipWriter.CreateHeader(header)
			if err != nil {
				file.Close()
				fmt.Printf("\n写入 zip 条目失败: %s\n%s\n", relPath, err)
				continue
			}
			_, err = io.Copy(writer, file)
			file.Close()
			if err != nil {
				fmt.Printf("\n写入失败: %s\n%s\n", relPath, err)
				continue
			}
			bar.Add(1)
		}
	} else {
		zipWriter := yekazip.NewWriter(zipFile)
		defer zipWriter.Close()

		for _, relPath := range filesToZip {
			fullPath := filepath.Join(basePath, relPath)
			file, err := os.Open(fullPath)
			if err != nil {
				fmt.Printf("\n无法打开: %s\n%s\n", relPath, err)
				continue
			}

			info, _ := file.Stat()
			header, err := yekazip.FileInfoHeader(info)
			if err != nil {
				file.Close()
				fmt.Printf("\n获取文件头失败: %s\n%s\n", relPath, err)
				continue
			}

			gbkNameBytes, err := Utf8ToGbk(relPath)
			if err != nil {
				file.Close()
				fmt.Printf("\n文件名转GBK失败: %s\n%s\n", relPath, err)
				continue
			}
			header.Name = string(gbkNameBytes)

			header.Method = zip.Deflate

			header.SetPassword(password)
			header.SetEncryptionMethod(yekazip.AES256Encryption)

			writer, err := zipWriter.Encrypt(header.Name, password, yekazip.AES256Encryption)
			if err != nil {
				file.Close()
				fmt.Printf("\n写入 zip 条目失败: %s\n%s\n", relPath, err)
				continue
			}

			_, err = io.Copy(writer, file)
			file.Close()
			if err != nil {
				fmt.Printf("\n写入失败: %s\n%s\n", relPath, err)
				continue
			}
			bar.Add(1)
		}
	}

	return nil
}

func collectFiles(basePath string, ign *ignore.GitIgnore) ([]string, error) {
	var filesToZip []string
	err := filepath.Walk(basePath, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		relPath, err := filepath.Rel(basePath, path)
		if err != nil {
			return err
		}

		if ign != nil && ign.MatchesPath(relPath) {
			if info.IsDir() {
				return filepath.SkipDir
			}
			return nil
		}
		if !info.IsDir() {
			filesToZip = append(filesToZip, relPath)
		}
		return nil
	})
	return filesToZip, err
}

func loadIgnore(srcDir string, ignoreFlag string) *ignore.GitIgnore {
	var patterns []string
	ignorePath := filepath.Join(srcDir, ".zipignore")
	if _, err := os.Stat(ignorePath); err == nil {
		lines, _ := os.ReadFile(ignorePath)
		fmt.Println("使用 .zipignore 忽略规则")
		for _, line := range strings.Split(string(lines), "\n") {
			if strings.TrimSpace(line) != "" {
				patterns = append(patterns, line)
			}
		}
	}

	if ignoreFlag != "" {
		cmdRules := strings.Split(ignoreFlag, ",")
		for _, rule := range cmdRules {
			rule = strings.TrimSpace(rule)
			if rule != "" {
				patterns = append(patterns, rule)
			}
		}
	}
	if len(patterns) > 0 {
		return ignore.CompileIgnoreLines(patterns...)
	}
	return nil
}

func processArgs(args []string) (string, string) {
	srcDir := args[0]
	if len(args) >= 2 {
		return srcDir, args[1]
	}
	folderName := filepath.Base(filepath.Clean(srcDir))
	outputZip := folderName + ".zip"
	fmt.Println("未指定输出文件名，使用默认:", outputZip)
	return srcDir, outputZip
}

func printUsage() {
	fmt.Println(`用法: zip [--ignore=规则] [--password=密码] <源文件夹> [输出zip路径]
=====================================
--ignore 选填，逗号分隔忽略规则，例如: "*.log,tmp/"
--password 选填，压缩密码
<源文件夹> 必填, 要压缩的文件路径
<输出zip路径> 选填, 默认使用文件夹名称.zip
=====================================
例子:
  zip ../myfolder
  zip ../myfolder ../myfolder.zip
  zip --ignore="target,dist" --password=123456 ../myfolder ../myfolder.zip`)
}

// Utf8ToGbk 将UTF-8字符串转为GBK编码
func Utf8ToGbk(s string) ([]byte, error) {
	reader := transform.NewReader(strings.NewReader(s), simplifiedchinese.GBK.NewEncoder())
	return io.ReadAll(reader)
}
